Use `narwhals` as dataframe-agnostic backend #7964

jessegrabowski · 2025-11-17T00:46:37Z

Description

Continues/completes #7462. I didn't have permission to push into that PR, so I'm opening this one.

The purpose of this PR is to use narwhals as a one-stop shop dataframe backend. Currently, we use pandas in data.py and pytensorf.py to allow users to pass dataframe objects into pm.Data and pt.as_tensor, respectively. I add a narwhals compatibility layer between the input and the pymc model to allow the user to bring his data in any form that narwhals supports, provided we register the libraries.

(If we could eliminate registration that would also be great, but I wasn't clever enough to figure out the multiple dispatch using only narwhals as a dependency. Maybe @MarcoGorelli could help 👉 👈 )

Some notes:

Since generalized dataframes don't have a notion of an index, we don't look at the index to find the labels for the left-most dimension provided to pm.Data. Instead, we look for a column matching that dimension name. If it is found, it is used as labels, and excluded from the values.
Narwhals has a lazy API via nw.LazyFrame and nw.LazySeries. I don't think we can do anything with those at the modeling level (maybe in the future with minibatching?). For now, I'm just calling .collect() on them to make them eager.
As mentioned above, we don't get all of narwhals for free yet -- the PR as it currently stands forces us to register each backend library we want to support. I don't think this is so bad, because it forces us to write tests for say DuckDB if someone comes along and really wants that. But it's a bit ugly. I added a dask.dataframe backend as an example of how we could extend things.

As of this PR, pandas could be made an optional or dev-only dependency for us. I didn't do it right away because I wanted to take people's temperature on the idea.

Related Issue

Closes ENH: Replace pandas dependence/use with narwhals #7462 Closes Add support for polars dataframes and series #7463
Related to #

Checklist

Checked that the pre-commit linting/style checks pass
Included tests that prove the fix is effective or that the new feature works
Added necessary documentation (docstrings and/or example notebooks)
If you are a pro: each commit corresponds to a relevant logical change

Type of change

Copilot

Pull Request Overview

This PR introduces narwhals as a dataframe-agnostic backend for PyMC, enabling support for multiple dataframe libraries (pandas, polars, dask.dataframe) through a unified interface. The implementation uses a registration mechanism to handle different dataframe backends dynamically.

Key changes:

Replaces direct pandas usage with narwhals compatibility layer
Implements singledispatch-based coordinate determination for different data types
Adds support for polars and dask.dataframe in addition to pandas

Reviewed Changes

Copilot reviewed 11 out of 12 changed files in this pull request and generated 11 comments.

Show a summary per file

File	Description
pymc/pytensorf.py	Implements dataframe backend registration and tensor conversion using narwhals
pymc/data.py	Refactors coordinate determination with singledispatch and narwhals integration
tests/test_pytensorf.py	Extends tests to cover pandas, polars, and dask.dataframe
tests/test_data.py	Adds polars-specific tests for series and dataframe coordinate inference
requirements.txt	Adds narwhals>=2.11.0 as a core dependency
requirements-dev.txt	Adds narwhals>=2.11.0 to development dependencies
conda-envs/*.yml	Adds narwhals>=2.11.0 to all conda environment files

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

pymc/pytensorf.py

pymc/data.py

requirements-dev.txt

pymc/data.py

tests/test_data.py

pymc/data.py

pymc/pytensorf.py

codecov · 2025-11-17T01:10:40Z

Codecov Report

❌ Patch coverage is 79.34783% with 19 lines in your changes missing coverage. Please review.
✅ Project coverage is 91.45%. Comparing base (f69e159) to head (ff7af55).

Files with missing lines	Patch %	Lines
pymc/data.py	75.71%	17 Missing ⚠️
pymc/pytensorf.py	90.90%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #7964      +/-   ##
==========================================
- Coverage   91.49%   91.45%   -0.04%     
==========================================
  Files         116      116              
  Lines       18962    19016      +54     
==========================================
+ Hits        17349    17391      +42     
- Misses       1613     1625      +12

Files with missing lines	Coverage Δ
pymc/pytensorf.py	`88.11% <90.90%> (+<0.01%)`	⬆️
pymc/data.py	`82.45% <75.71%> (-2.62%)`	⬇️

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

juanitorduz · 2025-11-19T17:00:40Z

I love this! I will be reviewing this soon :)

juanitorduz · 2025-11-20T08:53:28Z

pymc/data.py

+    return coords, _handle_none_dims(dims, value.ndim)
+
+
+def _dataframe_agnostic_coords(


Maybe we can add a simple test for this function?

I added a test for this. The new test revealed that my code for inferring an "index column" in the dataframe based on the left-most dim was broken.

I was trying to do was to allow for an "index column" in backends with no index (Polars). So if your dims were (date, country) and you passed in a polars dataframe with columns [USA, Chile, Mozambique, date], it would drop the "date" column, use it as a coord, then pass in the rest as the data to pm.Data. But this ended up being something of an API break, and I didn't like that we ended up mutating the user inputs silently in the background. I removed it, so only a pandas dataframe will give you index coords when infer_coords_and_dims=True

Frankly though I kind of hate this feature, it's too magical. Just force people to write out their coords.

juanitorduz

LGTM. Just a small comment / request

…ames

jessegrabowski requested review from cluhmann, Copilot, fonnesbeck, juanitorduz and ricardoV94 November 17, 2025 00:46

jessegrabowski added enhancements model labels Nov 17, 2025

github-actions bot added feature request hackathon Suitable for hackathon labels Nov 17, 2025

Copilot started reviewing on behalf of jessegrabowski November 17, 2025 00:47 View session

Copilot finished reviewing on behalf of jessegrabowski November 17, 2025 00:49

Copilot AI reviewed Nov 17, 2025

View reviewed changes

jessegrabowski force-pushed the narwhals-data branch from 72cfa96 to 9bcd9de Compare November 17, 2025 00:51

jessegrabowski force-pushed the narwhals-data branch from 9501ac3 to 48dcc85 Compare November 17, 2025 02:59

juanitorduz reviewed Nov 20, 2025

View reviewed changes

juanitorduz mentioned this pull request Nov 20, 2025

Support additional DataFrame libraries pymc-labs/pymc-marketing#1691

Open

edschofield and others added 10 commits November 24, 2025 21:17

Generalize data input to support Narwhals-compatible Series or DataFr…

f7e4b20

…ames

Narwhals compatibility layer in pm.Data

0bcb631

fix typehint

edc1fe7

mypy T___T

86a7b05

Helpful robot

d32abc8

Remove dask series test (no narwhals support)

8a10ba3

Respond to feedback

f65ad77

mypy

a5da840

Small bugfixes and test

3553f00

code cleanup + mypy

ff7af55

jessegrabowski force-pushed the narwhals-data branch from a8b264e to ff7af55 Compare November 25, 2025 03:27

jessegrabowski requested a review from juanitorduz November 27, 2025 05:15

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Use `narwhals` as dataframe-agnostic backend #7964

Use `narwhals` as dataframe-agnostic backend #7964

Uh oh!

jessegrabowski commented Nov 17, 2025

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov bot commented Nov 17, 2025 •

edited

Loading

Uh oh!

juanitorduz commented Nov 19, 2025 •

edited

Loading

Uh oh!

juanitorduz Nov 20, 2025

Uh oh!

jessegrabowski Nov 25, 2025

Uh oh!

juanitorduz left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return coords, _handle_none_dims(dims, value.ndim)


		def _dataframe_agnostic_coords(

Use narwhals as dataframe-agnostic backend #7964

Are you sure you want to change the base?

Use narwhals as dataframe-agnostic backend #7964

Uh oh!

Conversation

jessegrabowski commented Nov 17, 2025

Description

Related Issue

Checklist

Type of change

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov bot commented Nov 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

juanitorduz commented Nov 19, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

juanitorduz Nov 20, 2025

Choose a reason for hiding this comment

Uh oh!

jessegrabowski Nov 25, 2025

Choose a reason for hiding this comment

Uh oh!

juanitorduz left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Use `narwhals` as dataframe-agnostic backend #7964

Use `narwhals` as dataframe-agnostic backend #7964

codecov bot commented Nov 17, 2025 •

edited

Loading

juanitorduz commented Nov 19, 2025 •

edited

Loading